'''本脚本用于测试图片识别正确率'''
import os
import sys
from PIL import Image
from io import BytesIO
import re
import Levenshtein

BASE_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_PATH)

from model import  OcrHandle
from config import dbnet_max_size

ocrhandle = OcrHandle()

# 接收药品图片的字节串和处方名，返回匹配结果
def OCR(raw_img, recipe):
    img = Image.open(BytesIO(raw_img))
    if hasattr(img, '_getexif') and img._getexif() is not None:
        orientation = 274
        exif = dict(img._getexif().items())
        if orientation not in exif:
            exif[orientation] = 0
        if exif[orientation] == 3:
            img = img.rotate(180, expand=True)
        elif exif[orientation] == 6:
            img = img.rotate(270, expand=True)
        elif exif[orientation] == 8:
            img = img.rotate(90, expand=True)
    img = img.convert("RGB")

    compress_size = 960
    short_size = compress_size
    short_size = 32 * (short_size//32)

    img_w, img_h = img.size
    if max(img_w, img_h) * (short_size * 1.0 / min(img_w, img_h)) > dbnet_max_size:
        print("图片reize后长边过长，请调整短边尺寸")
        return None

    res = ocrhandle.text_predict(img,short_size)

    pattern = re.compile("'[0-9]+、 (.*)'")
    textList = pattern.findall(str(res))
    simtext = []
    for text in textList:
        similarity = Levenshtein.ratio(recipe, text)
        if similarity > 0.3:
            simtext.append(text)
        if (similarity > 0.7):
            if len(text) == len(recipe) or similarity > 0.85:
                return True
    print(simtext)  # 显示识别出的与处方相似的文本，用来调整字符串匹配的阈值
    return False


img_cnt = 0
correct_cnt = 0

print("测试开始")
print("注：只显示识别不正确的样例")

dir_names = []
for dir_name in os.listdir("F:/软B/测试"):
    dir_names.append(dir_name)

for dir_name in dir_names:
    for img_name in os.listdir("F:/软B/测试/%s" %dir_name):
        f = open("F:/软B/测试/%s/%s" %(dir_name, img_name), "rb")
        raw_img = f.read()
        f.close()
        img_cnt += 1
        OCR_result = OCR(raw_img, dir_name)
        if OCR_result is None or not OCR_result:
            print("识别不正确，图片路径：F:/软B/测试/%s/%s" %(dir_name, img_name))
        else:
            correct_cnt += 1

precision = correct_cnt / img_cnt
print("测试结束")
print("识别图片总数：", img_cnt)
print("识别正确图片数：", correct_cnt)
print("准确率：", correct_cnt/img_cnt)